package com.it.smn.divide.datasource;

import com.it.smn.divide.loadbalance.LoadBalance;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.util.Assert;

import javax.sql.DataSource;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;

/**
 * @author sunmannian
 * @ClassName com.it.smn.divide.datasource.DynamicDataSourceRouter
 * @Date 2020/8/27 15:17
 * @Version V1.0
 * Description  <请输入一句话描述该类功能>
 * Copyright (C) 北京中电兴发科技有限公司  since 2020/8/27
 */
public class DynamicDataSourceRouter implements InitializingBean,
        ApplicationContextAware {

    private List<DataSource> writeDataSources;

    private List<DataSource> readDataSources;

    private Map<String, DataSource> writeDataSourcesMap;

    private Map<String, DataSource> readDataSourcesMap;

    private ApplicationContext applicationContext;

    private LoadBalance loadBalance;

    /**
     * 后续可实现负载均衡
     */
    public LoadBalance getLoadBalance() {
        return loadBalance;
    }

    public void setLoadBalance(LoadBalance loadBalance) {
        this.loadBalance = loadBalance;
    }


    @Override
    public void afterPropertiesSet() {
        Assert.notNull(this.loadBalance, "loadBalance must not be null");
        Map<String, WritableDataSource> writeDataSourceMap = applicationContext.getBeansOfType(WritableDataSource.class, false, true);
        Map<String, ReadableDataSource> readDataSourceMap = applicationContext.getBeansOfType(ReadableDataSource.class, false, true);
        writeDataSources = new CopyOnWriteArrayList<>(writeDataSourceMap.values());
        readDataSources = new CopyOnWriteArrayList<>(readDataSourceMap.values());
        writeDataSourcesMap = new ConcurrentHashMap<>(writeDataSourceMap);
        readDataSourcesMap = new ConcurrentHashMap<>(readDataSourceMap);
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    public DataSource getReadDataSource() {
        DataSource dataSource = loadBalance.getLoadBalanceStrategy().select(readDataSources);
        Assert.notNull(dataSource, "datasource must not be null");
        return dataSource;
    }

    public DataSource getReadDataSourceByName(String name) {
        DataSource dataSource = readDataSourcesMap.get(name);
        Assert.notNull(dataSource, "datasource must not be null");
        return dataSource;
    }

    public DataSource getWriteDataSource() {
        DataSource dataSource = loadBalance.getLoadBalanceStrategy().select(writeDataSources);
        Assert.notNull(dataSource, "datasource must not be null");
        return dataSource;
    }

    public DataSource getWriteDataSourceByName(String name) {
        DataSource dataSource = writeDataSourcesMap.get(name);
        Assert.notNull(dataSource, "datasource must not be null");
        return dataSource;
    }

    public List<DataSource> getWriteDataSources() {
        return writeDataSources;
    }

    public void setWriteDataSources(List<DataSource> writeDataSources) {
        this.writeDataSources = writeDataSources;
    }

    public List<DataSource> getReadDataSources() {
        return readDataSources;
    }

    public void setReadDataSources(List<DataSource> readDataSources) {
        this.readDataSources = readDataSources;
    }

    public Map<String, DataSource> getWriteDataSourcesMap() {
        return writeDataSourcesMap;
    }

    public void setWriteDataSourcesMap(Map<String, DataSource> writeDataSourcesMap) {
        this.writeDataSourcesMap = writeDataSourcesMap;
    }

    public Map<String, DataSource> getReadDataSourcesMap() {
        return readDataSourcesMap;
    }

    public void setReadDataSourcesMap(Map<String, DataSource> readDataSourcesMap) {
        this.readDataSourcesMap = readDataSourcesMap;
    }

    public ApplicationContext getApplicationContext() {
        return applicationContext;
    }


}
